#!/usr/bin/env python3
# encoding: utf-8
# @data:2023/07/11
# @author:aiden
# 自动踢球
import math
import rospy
import signal
from std_msgs.msg import Float64, String
from ainex_sdk import pid, misc, common
from ainex_kinematics.gait_manager import GaitManager
from ainex_example.color_common import Common
from ainex_example.pid_track import PIDTrack
from ainex_example.approach_object import ApproachObject
from ainex_interfaces.srv import SetString
from ainex_interfaces.msg import ObjectsInfo, ColorDetect, ColorsDetect

class KickBallNode(Common):
    # 左右踢球的动作名
    left_shot_action_name = 'new_left_shot'
    right_shot_action_name = 'new_right_shot'
    forward_step = 'forward_one_step'  #前进一步
    back_step = 'back_step'        #后退一步
    move_left = 'move_left'        #左移
    move_right = 'move_right'      #右移
    

    door_roi = [0, 1, 0.1, 0.9]
    ball_roi = [0, 1, 0, 1] 

    # 图像处理大小  
    image_process_size = [160, 120]

    door_x_stop = 250/640
    ball_y_stop = 350/480
    ball_x_stop = 420/640
    
    kick_y_stop = 285/480
    kick_x_stop = 200/640

    def __init__(self, name):
        rospy.init_node(name)
        self.name = name
        self.gait_manager = GaitManager()
        self.go_gait_param = self.gait_manager.get_gait_param()
        self.go_gait_param['body_height'] = 0.025
        self.go_gait_param['step_height'] = 0.02
        self.go_gait_param['hip_pitch_offset'] = 12
        self.go_gait_param['z_swap_amplitude'] = 0.006
        self.go_gait_param['init_roll_offset'] = -3
        self.go_dsp = [400, 0.2, 0.03]
        self.go_arm_swap = 30

        self.door_dsp = [300, 0.2, 0.025]

        self.running = True
        self.left = False
        self.right = False
        # 存储检测结果
        self.objects_info = []
        self.count_miss = 0
        # 找球相关标志位
        self.start_index = 0
        self.start_find_ball = False
        # 初始化头部位置
        self.head_pan_init = 500  # 左右舵机的初始值
        self.head_tilt_init = 300 # 上下舵机的初始值
        self.head_time_stamp = rospy.get_time()
        # 初始化父类 
        super().__init__(name, self.head_pan_init, self.head_tilt_init)
        self.calib_config = common.get_yaml_data('/home/ubuntu/ros_ws/src/ainex_example/config/calib.yaml')
        # 存放检测结果
        self.rl_dis = None
        self.ud_dis = None
        # 初始化接近控制
        self.approach_object = ApproachObject(self.gait_manager, step_mode=0)
        self.approach_object.update_gait(dsp=[400, 0.2, 0.025])
        self.approach_object.update_gait(walking_param = self.go_gait_param)
        self.approach_object.update_stop_count(1)
        self.approach_object.update_gait_range(x_range=[-0.013, 0.013])
        self.approach_object.update_approach_stop_value(30, 0, 3)
        signal.signal(signal.SIGINT, self.shutdown)

        self.current_state = "ball"    #当前状态
        #状态列表
        self.state = {'door':   [[500, 400], ['red', self.door_roi, self.image_process_size, self.set_door_color], False],
                      'ball': [[500, 400], ['blue', self.ball_roi, self.image_process_size, self.set_ball_color], False],
                      'kick': [[500, 260], ['blue', self.ball_roi, self.image_process_size, self.set_ball_color], False]}


        # 头部的pid追踪
        self.head_pan_range = [125, 875]  # 左右转动限制在这个范围， 125为右
        self.head_tilt_range = [260, 500]  # 上下限制在这个范围， 250为下
        self.pid_rl = pid.PID(0.1, 0.0, 0.001)
        self.pid_ud = pid.PID(0.1, 0.0, 0.001)
        self.head_pan_init = 500  # 左右舵机的初始值
        self.head_tilt_init = 400 # 上下舵机的初始值
        self.rl_track = PIDTrack(self.pid_rl, self.head_pan_range, self.head_pan_init)
        self.ud_track = PIDTrack(self.pid_ud, self.head_tilt_range, self.head_tilt_init)
        
        # 躯体的追踪参数
        self.yaw_stop = 40  # 躯体停止转动时头部左右舵机脉宽值和中位500的差值
        self.x_stop = 340
        # 找球时头部经过的5个位置，左右舵机，上下舵机，时间ms
        # left_down, left_up, center_up, right_up, right_down
        self.find_ball_position = [[650, 300, 1000], 
                                   [650, 500, 1000], 
                                   [500, 500, 1000], 
                                   [350, 500, 1000],
                                   [350, 300, 1000]
                                   ]

        # 订阅颜色识别结果
        rospy.Subscriber('/object/pixel_coords', ObjectsInfo, self.get_color_callback)
        rospy.Service('~set_color', SetString, self.set_color_srv_callback)  # 设置颜色
        # 初始化站姿  
        self.motion_manager.run_action('walk_ready')
        # 是否自动开始
        if rospy.get_param('~start', True):
            # 通知颜色识别准备，此时只显示摄像头原画
            self.enter_func(None)
            self.start_srv_callback(None)
            common.loginfo('start ombination')
    # 退出回调 
    def shutdown(self, signum, frame):
        with self.lock:
            self.running = False
            common.loginfo('%s shutdown' % self.name)
    
    def set_ball_color(self, color, roi, image_process_size):
        # 设置追踪颜色
        ball_param = ColorDetect()
        ball_param.color_name = color
        ball_param.detect_type = 'circle'
        ball_param.use_name = True
        ball_param.image_process_size = self.image_process_size
        ball_param.roi.y_min = int(roi[0] * image_process_size[1])
        ball_param.roi.y_max = int(roi[1] * image_process_size[1])
        ball_param.roi.x_min = int(roi[2] * image_process_size[0])
        ball_param.roi.x_max = int(roi[3] * image_process_size[0])
        ball_param.min_area = 20
        ball_param.max_area = self.image_process_size[0]*self.image_process_size[1]

        return ball_param
    
    def set_door_color(self, color, roi, image_process_size):
        # 设置门颜色
        door_param = ColorDetect()
        door_param.color_name = color
        door_param.detect_type = 'rect'
        door_param.use_name = True
        door_param.image_process_size = self.image_process_size
        door_param.roi.y_min = int(roi[0] * image_process_size[1])
        door_param.roi.y_max = int(roi[1] * image_process_size[1])
        door_param.roi.x_min = int(roi[2] * image_process_size[0])
        door_param.roi.x_max = int(roi[3] * image_process_size[0])
        door_param.min_area = 20
        door_param.max_area = self.image_process_size[0]*self.image_process_size[1]

        return door_param
    
    def set_color_srv_callback(self, msg):
        # 设置颜色
        door_param = self.set_stairs_color(self.state['door'][1][0])
        ball_param = self.set_visual_patrol_color(self.state['ball'][1][0])

        self.detect_pub.publish([door_param, ball_param])
        common.loginfo('%s set_color' % self.name)
        
        return [True, 'set_color']
    

    def state_init(self, current_state):
        # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param1 = self.state[current_state][1][3](self.state[current_state][1][0], self.state[current_state][1][1], self.state[current_state][1][2])
            param2 = self.state['ball'][1][3](self.state['ball'][1][0], self.state['ball'][1][1], self.state['ball'][1][2])
            self.detect_pub.publish([param1, param2])  # 颜色检测设置
            common.loginfo(current_state + ' init')

    def get_color_callback(self, msg):
        # 获取颜色识别结果
        self.objects_info = msg.data
        # 获取识别结果，根据结果计算距离

    # 躯体追踪
    def body_track_process(self, ball_data):                                                
        # 左右根据头部左右舵机位置值进行调整
        if ball_data is not None:
            print(ball_data.y)
            if ball_data.y < self.ball_y_stop * ball_data.height:  
                self.gait_manager.set_step(self.go_dsp, 0.02, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            # elif ball_data.x - self.ball_x_stop * ball_data.width > 20:
            #     self.gait_manager.set_step(self.go_dsp, 0, -0.012, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            # elif ball_data.x - self.ball_x_stop * ball_data.width < -20:
            #     self.gait_manager.set_step(self.go_dsp, 0, 0.012, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)  
            else:
                return True        
            rospy.sleep(0.1)  
        return False  

    def exit_door(self, door_data, ball_data):
        if door_data is not None and ball_data is not None:
            print(door_data.x, ball_data.x)
            print(door_data)
            if door_data.x + 120 > 290:
                self.gait_manager.set_step(self.go_dsp, 0, 0, -4, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            elif ball_data.x -280 > 10:
                self.gait_manager.set_step(self.go_dsp, 0, -0.01, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            elif ball_data.x -280 < -10:
                self.gait_manager.set_step(self.go_dsp, 0, 0.01, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            else:
                rospy.sleep(1)
                return True 
            rospy.sleep(0.1)
        elif ball_data is None:
            self.gait_manager.set_step(self.go_dsp, -0.01, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
        return False
    
    def kick_ball(self, ball_data):
        if ball_data is not None:
            print(ball_data.x, ball_data.y)
            print(ball_data.width, ball_data.height)
            if  ball_data.y < self.kick_y_stop * ball_data.height:   
                self.gait_manager.set_step(self.go_dsp, 0.008, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0) 
            elif  ball_data.x - self.kick_x_stop * ball_data.width < -10:
                self.gait_manager.set_step(self.go_dsp, 0, 0.005, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            elif ball_data.x - self.kick_x_stop * ball_data.width > 10:
                self.gait_manager.set_step(self.go_dsp, 0, -0.005, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)   
            else:
                self.gait_manager.disable()
                self.motion_manager.run_action(self.left_shot_action_name)
                self.running = False
            rospy.sleep(0.1)
            #print(ball_data)
            #if ball_data.y < 20:

    # 找球过程 
    def find_ball_process(self):
        # 根据预设的扫描点进行找球
        if rospy.get_time() > self.head_time_stamp:
            if self.start_index > len(self.find_ball_position) - 1:
                self.start_index = 0
            rl_dis = self.find_ball_position[self.start_index][0]
            ud_dis = self.find_ball_position[self.start_index][1]
            self.rl_track.update_position(rl_dis)  # pid的输出值要跟着更新
            self.ud_track.update_position(ud_dis)
            self.motion_manager.set_servos_position(self.find_ball_position[self.start_index][2], [[23, rl_dis], [24, ud_dis]])

            self.gait_manager.move(2, 0, 0, 5)  # 右转
            self.head_time_stamp = rospy.get_time() + self.find_ball_position[self.start_index][2]/1000.0
            self.start_index += 1

    def run(self):
        while self.running:
            # 状态判断
            if self.start:
                if self.current_state == "door":
                    door_data = None
                    ball_data = None
                    for object_info in self.objects_info:
                        if object_info.type == 'rect':
                            door_data = object_info
                        if object_info.type == 'circle':
                            ball_data = object_info
                    if self.exit_door(door_data, ball_data):
                        self.current_state = "kick"
                
                elif self.current_state == "kick":
                    ball_data = None
                    for object_info in self.objects_info:
                        if object_info.type == 'circle':
                            ball_data = object_info
                    self.kick_ball(ball_data)

                elif self.current_state == "ball":
                    ball_data = None 
                    # 在检测到的对象中查找球
                    for object_info in self.objects_info:
                        if object_info.type == 'circle':
                            ball_data = object_info
                    # 如果识别到球就进行躯体追踪
                    if ball_data is not None:
                        if self.body_track_process(ball_data):
                            self.current_state = "door"
                        self.start_find_ball = False
                
                    rospy.sleep(0.01)
                self.state_init(self.current_state)

            else:
                rospy.sleep(0.01)

        self.init_action(self.head_pan_init, self.head_tilt_init)
        self.stop_srv_callback(None)
        rospy.signal_shutdown('shutdown')

if __name__ == "__main__":
    KickBallNode('kick_ball').run()
